from typing import Iterable, List, Union
import attr
from pyecharts import options as opts


@attr.s
class KlineBuilder:

    # init
    width: str = attr.ib('900px')
    height: str = attr.ib('800px')
    pos_left: str = attr.ib('5%')
    pos_right: str = attr.ib('1%')
    main_height: int = attr.ib(60)
    is_zoom: bool = attr.ib(False)
    graph_objs: dict = attr.ib(factory=dict)
    main: str = attr.ib('kline')
    minors: List[str] = attr.ib(factory=list)
    minors_weights: List[int] = attr.ib(factory=list)
    zoom_opts: dict = attr.ib(factory=dict)
    gap_main_minor: int = attr.ib(10)
    gap_minors: int = attr.ib(1)
    slider_height: int = attr.ib(5)

    def set_jupyter_server(self, port: int = 8888):
        """离线环境下，配置notebook，挂载echarts服务
        
        :param port: notebook的端口
        """
        from pyecharts.globals import CurrentConfig
        # OnlineHostType.NOTEBOOK_HOST 默认值为 http://localhost:8888/nbextensions/assets/
        setattr(CurrentConfig, 'ONLINE_HOST',
                f"http://localhost:{port}/nbextensions/assets/")
        return self

    def set_notebook(self):
        """设置当前的渲染模式为notebook
        
        pyecharts默认使用notebook渲染
        """
        from pyecharts.globals import CurrentConfig, NotebookType
        setattr(CurrentConfig, 'NOTEBOOK_TYPE', NotebookType.JUPYTER_NOTEBOOK)
        return self

    def set_jupyter_lab(self):
        """设置当前的渲染模式为jupyter lab
        """
        from pyecharts.globals import CurrentConfig, NotebookType
        setattr(CurrentConfig, 'NOTEBOOK_TYPE', NotebookType.JUPYTER_LAB)
        return self

    @classmethod
    def get_kline(cls, times, ohlc, series_name: str = 'kline', title: str = 'K线周期图表'):
        """生成k线图对象
        
        :param times: 时间标签数据
        :param ohlc: ohlc数据（二维数组），注意列顺序为开盘、收盘、最低、最高。
        :param title: K线图的标题
        """
        from pyecharts.charts import Kline

        kline = (
            Kline()
            .add_xaxis(xaxis_data=times)
            .add_yaxis(
                series_name=series_name,
                y_axis=ohlc,
                itemstyle_opts=opts.ItemStyleOpts(
                    color="#ef232a",
                    color0="#14b143",
                    border_color="#ef232a",
                    border_color0="#14b143",
                ),
                markpoint_opts=opts.MarkPointOpts(
                    data=[
                        opts.MarkPointItem(
                            type_="max", name="最大值", symbol='arrow', symbol_size=30),
                        opts.MarkPointItem(
                            type_="min", name="最小值", symbol='arrow', symbol_size=30),
                    ],
                    label_opts=opts.LabelOpts(position='inside'),
                ),
            )
            .set_global_opts(
                title_opts=opts.TitleOpts(title=title, pos_left="0"),
                xaxis_opts=opts.AxisOpts(
                    type_="category",
                    is_scale=True,
                    boundary_gap=False,
                    axisline_opts=opts.AxisLineOpts(is_on_zero=False),
                    splitline_opts=opts.SplitLineOpts(is_show=False),
                    split_number=20,
                    min_="dataMin",
                    max_="dataMax",
                ),
                yaxis_opts=opts.AxisOpts(
                    is_scale=True,
                    # splitline_opts=opts.SplitLineOpts(is_show=True),
                    splitarea_opts=opts.SplitAreaOpts(
                        is_show=True, areastyle_opts=opts.AreaStyleOpts(opacity=1)
                    ),
                ),
                tooltip_opts=opts.TooltipOpts(
                    trigger="axis",
                    axis_pointer_type="cross",
                    background_color="rgba(245, 245, 245, 0.8)",
                    border_width=1,
                    border_color="#ccc",
                    textstyle_opts=opts.TextStyleOpts(color="#000"),
                ),
                axispointer_opts=opts.AxisPointerOpts(
                    is_show=True,
                    link=[{"xAxisIndex": "all"}],
                    label=opts.LabelOpts(background_color="#777"),
                ),
            )
        )
        return kline

    def add_kline(self, times, ohlc, series_name: str = '', name: str = 'kline', title: str = 'K线周期图表'):
        """添加k线图对象
        
        :param times: 时间标签数据
        :param ohlc: ohlc数据（二维数组），注意列顺序为开盘、收盘、最低、最高。
        :param series_name: 数据序列的名称，用于内部的调用，以及作为数据标签的名称。如果没有定义的话，默认和`name`一样。
        :param name: 图表的名称，用于内部的调用
        :param title: K线图的标题
        """
        series_name = series_name or name
        self.graph_objs[name] = self.get_kline(
            times, ohlc, series_name=series_name, title=title)
        return self

    def add_zoom_slider(self, pos_top: str = "97%"):
        """添加缩放图形滑块
        
        :param post_top: 滑块举例整个画布顶部的距离
        """
        self.zoom_opts['slider'] = opts.DataZoomOpts(
            is_show=True, xaxis_index=[0], pos_top=pos_top, range_end=100
        )
        return self

    def add_zoom_inside(self):
        """增加图形内部缩放功能
        
        """
        self.zoom_opts['inside'] = opts.DataZoomOpts(
            is_show=False, type_='inside', xaxis_index=list(range(len(self.minors) + 1)),
            range_end=100,
        )
        return self

    def _set_zoom(self):
        """更新zoom配置"""
        self.graph_objs[self.main].options.update(
            dataZoom=[*self.zoom_opts.values()])
        return self

    def with_zoom(self):
        """设置图形需要缩放功能
        
        设置此项后，会默认为图形增加内部缩放和缩放滑块的功能
        """
        self.is_zoom = True
        return self

    def markarea_xaxis(self,
                       begin, end,
                       name: str = '', series_name: str = '',
                       color: str = "#DCA3A2", opacity: float = 0.5,
                       show_label: bool = False):
        """垂直X轴标记一片面积
        
        :param begin: x轴起始位置
        :param end: x轴终止位置
        :param series_name: 数据序列的名称，用于内部的调用，以及作为数据标签的名称。如果没有定义的话，默认和`name`一样。
        :param name: 图表的名称，用于内部的调用
        :param color: 标记区域的颜色
        :param opacity: 标记区域的透明度
        :param show_label: 是否展示图标
        """
        series_name = series_name or name
        markarea_opts = opts.MarkAreaOpts(
            is_silent=False,
            data=[opts.MarkAreaItem(
                name=series_name,
                x=(begin, end),
                label_opts=opts.LabelOpts(is_show=show_label),
                itemstyle_opts=opts.ItemStyleOpts(
                    color=color, opacity=opacity),
            )]
        )
        for s in self.graph_objs[name].options.get('series'):
            s.update(markArea=markarea_opts)
        return self

    def markarea_xaxis_many(self,
                            rngs,
                            name: str = '', series_name: str = '',
                            color: str = "#DCA3A2", opacity: float = 0.5,
                            show_label: bool = False):
        """垂直X轴标记多片面积
        
        :param rngs: 一系列的元组，元组的第一个元素是x轴起始位置，第二个元素是x周的终止位置
        :param series_name: 数据序列的名称，用于内部的调用，以及作为数据标签的名称。如果没有定义的话，默认和`name`一样。
        :param name: 图表的名称，用于内部的调用
        :param color: 标记区域的颜色
        :param opacity: 标记区域的透明度
        :param show_label: 是否展示图标
        """
        series_name = series_name or name
        markarea_opts = opts.MarkAreaOpts(
            is_silent=False,
            data=[opts.MarkAreaItem(
                name=series_name,
                x=(begin, end),
                label_opts=opts.LabelOpts(is_show=show_label),
                itemstyle_opts=opts.ItemStyleOpts(
                    color=color, opacity=opacity),
            ) for begin, end in rngs]
        )
        for s in self.graph_objs[name].options.get('series'):
            s.update(markArea=markarea_opts)
        return self

    def add_main_title(self, title: str = 'K线图'):
        """设置主图的标题"""
        title_opts = opts.TitleOpts(title=title, pos_left="0")
        self.graph_objs[self.main].options.update(title=title_opts)
        return self

    @classmethod
    def get_bar(cls, times, values, series_name: str = 'Volume', colors: Union[Iterable[str], str, None] = None):
        """创建柱状图对象
        
        :param times: 时间标签数据
        :param values: 数据
        :param series_name: 数据序列的名称，用于内部的调用，以及作为数据标签的名称。如果没有定义的话，默认和`name`一样。
        :param colors: 柱子的颜色，当为str时会统一一个颜色；当是一系列str时，为每个柱子设置不同的颜色；None的时候分配默认的颜色
        """
        from pyecharts.commons.utils import JsCode
        from pyecharts.charts import Bar

        if isinstance(colors, Iterable):
            color_func = JsCode(
                """
                function(params) {{
                    var colors = {};
                    return colors[params.dataIndex]
                }}
                """.format(list(colors)))
            item_opt = opts.ItemStyleOpts(color=color_func)
        elif isinstance(colors, str):
            item_opt = opts.ItemStyleOpts(color=colors)
        else:
            item_opt = None
        bar = (
            Bar()
            .add_xaxis(xaxis_data=times)
            .add_yaxis(
                series_name=series_name,
                y_axis=values,
                label_opts=opts.LabelOpts(is_show=False),
                itemstyle_opts=item_opt,
            )
            .set_global_opts(
                xaxis_opts=opts.AxisOpts(
                    type_="category",
                    axislabel_opts=opts.LabelOpts(is_show=False),
                ),
                legend_opts=opts.LegendOpts(is_show=False),
            )
        )

        return bar

    def add_bar(self, times, values, series_name: str = '',  name='volume', colors=None):
        """添加柱状图对象
        
        :param times: 时间标签数据
        :param values: 数据
        :param name: 图表的名称，用于内部调用
        :param series_name: 数据序列的名称，用于内部的调用，以及作为数据标签的名称。如果没有定义的话，默认和`name`一样。
        :param colors: 柱子的颜色，当为str时会统一一个颜色；当是一系列str时，为每个柱子设置不同的颜色；None的时候分配默认的颜色
        """
        series_name = series_name or name
        self.graph_objs[name] = self.get_bar(
            times, values, series_name=series_name, colors=colors)
        return self

    @classmethod
    def get_colors_by_sign(cls, values, pos='#ef232a', neg='#14b143'):
        """根据序列中每个元素的符号，生成颜色序列
        
        方便为柱状图根据涨跌上色
        
        :param values: 数据序列
        :param pos: 符号为正时对应的颜色
        :param neg: 符号为负时对应的颜色
        """
        import numpy as np
        return np.where(values > 0, pos, neg)

    @classmethod
    def get_line(cls, times, values, series_name: str = '', is_smooth: bool = True, color: str = ''):
        """创建折线图对象
        
        :param times: 时间标签数据
        :param values: 数据
        :param series_name: 数据序列的名称，用于内部的调用，以及作为数据标签的名称。
        :param is_smooth: 是否平滑
        :param color: 线的颜色
        """
        from pyecharts.charts import Line

        item_opt = opts.ItemStyleOpts(
            color=color) if color else None

        kline_line = (
            Line()
            .add_xaxis(xaxis_data=times)
            .add_yaxis(
                series_name=series_name,
                y_axis=values,
                is_smooth=is_smooth,
                linestyle_opts=opts.LineStyleOpts(opacity=0.5),
                label_opts=opts.LabelOpts(is_show=False),
                itemstyle_opts=item_opt
            )
            .set_global_opts(
                xaxis_opts=opts.AxisOpts(
                    type_="category",
                    axislabel_opts=opts.LabelOpts(is_show=False),
                ),
                yaxis_opts=opts.AxisOpts(
                    split_number=3,
                    axisline_opts=opts.AxisLineOpts(is_on_zero=False),
                    axistick_opts=opts.AxisTickOpts(is_show=False),
                    splitline_opts=opts.SplitLineOpts(is_show=False),
                    axislabel_opts=opts.LabelOpts(is_show=True),
                ),
            )
        )
        return kline_line

    def add_line(self, times, values, series_name: str = '', color: str = '', name: str = 'line'):
        """添加折线图对象
        
        :param times: 时间标签数据
        :param values: 数据
        :param name: 图表的名称，用于内部调用
        :param series_name: 数据序列的名称，用于内部的调用，以及作为数据标签的名称。
        :param is_smooth: 是否平滑
        :param color: 线的颜色
        """
        series_name = series_name or name
        self.graph_objs[name] = self.get_line(times, values,
                                              series_name=series_name,
                                              color=color)
        return self

    def overlap(self, first: str, second: str):
        """指定第二个图叠加到第一个图上
        
        :param first: 第一幅图表的名称
        :param second: 第二幅图表的名称
        """
        f_g = self.graph_objs[first]
        s_g = self.graph_objs[second]
        self.graph_objs[first] = f_g.overlap(s_g)
        return self

    def set_main(self, name: str = 'kline'):
        """设置主图
        
        默认为kline
        
        :param name: 图表的名称
        """
        self.main = name
        return self

    def add_minor(self, name: str):
        """设置副图
        """
        if name in self.minors:
            raise ValueError(f'{name} already in minors')
        self.minors.append(name)
        return self

    def add_minors(self, *name: str):
        """指定多个图表副图
        """
        for n in name:
            self.add_minor(n)
        return self

    def set_figsize(self, width: str = '900px', height: str = '800px'):
        """设置整个图的大小
        """
        self.width = width
        self.height = height
        return self

    def set_minors_weights(self, weights: List[int]):
        """设置副图大小的权重
        
        不设置的时候就是等权重的，一样大
        """
        if len(weights) != len(self.minors_weights):
            raise ValueError('weights must be same length with minors!')
        self.minors_weights = weights
        return self

    def add_echart(self, obj, name: str):
        """添加一个echart图形对象
        
        方便添加自定义的echart图表
        
        :param name: 图表的名称，用于内部调用
        """
        self.graph_objs[name] = obj
        return self

    def build(self):
        """构建整个图表
        
        :return: 返回图表对象，可用`xxx.render_notebook()`在notebook或jupyter lab中渲染图形
        """
        from pyecharts import options as opts
        from pyecharts.charts import Grid

        # set zoom
        if self.is_zoom:
            self.add_zoom_inside()
            self.add_zoom_slider()
        if self.zoom_opts:
            self._set_zoom()

        grid_chart = Grid(init_opts=opts.InitOpts(
            width=self.width, height=self.height))
        if self.minors:
            grid_chart.add(
                self.graph_objs[self.main],
                grid_opts=opts.GridOpts(
                    pos_left=self.pos_left, pos_right=self.pos_right, height=f"{self.main_height}%"),
            )

            minor_height = 100 - self.main_height - self.gap_main_minor
            if self.zoom_opts.get('slider'):
                minor_height -= self.slider_height

            if not self.minors_weights:
                minor_num = len(self.minors)
                div_h = [minor_height // minor_num] * minor_num
            else:
                tot = sum(self.minors_weights)
                div_h = (int(w/tot * minor_height)
                         for w in self.minors_weights)

            for i, (h, minor) in enumerate(zip(div_h, self.minors)):
                pos_top = self.main_height + self.gap_main_minor + \
                    i * h + (i + 1) * self.gap_minors
                grid_chart.add(
                    self.graph_objs[minor],
                    grid_opts=opts.GridOpts(pos_left=self.pos_left,
                                            pos_right=self.pos_right,
                                            pos_top=f"{pos_top}%",
                                            height=f'{h}%'),
                )
        else:
            grid_chart.add(
                self.graph_objs[self.main],
                grid_opts=opts.GridOpts(
                    pos_left=self.pos_left, pos_right=self.pos_right, height=f"{100 - self.slider_height}%"),
            )
        return grid_chart
